from torchline.models import META_ARCH_REGISTRY

import torch
import torch.nn as nn


__all__ = [
    'SimpleCNN'
]

@META_ARCH_REGISTRY.register()
class SimpleCNN(nn.Module):
    def __init__(self, cfg):
        super(SimpleCNN, self).__init__()
        C = 32
        self.stem = ResBlock(2, C, 1)

        self.down1 = nn.Sequential(
            ResBlock(C, C, 1),
            ResBlock(C, 2*C, 2),
        )
        self.down2 = nn.Sequential(
            ResBlock(2*C, 2*C, 1),
            ResBlock(2*C, 4*C, 2)
        )
        self.down3 = nn.Sequential(
            ResBlock(4*C, 4*C, 1),
            ResBlock(4*C, 8*C, 2)
        )
        # self.down4 = nn.Sequential(
        #     ResBlock(8*C, 8*C, 1),
        #     ResBlock(8*C, 16*C, 2)
        # )
        self.pread_layer = nn.Linear(8*C*8, 1024, bias=False)
        self.init()

    def forward(self, x): # b*2*512
        bs = x.shape[0]
        out_stem = self.stem(x) # b*C*512
        out = torch.nn.MaxPool1d(2)(out_stem) # b*C*256
        out = self.down1(out)   # b*2C*128
        out = torch.nn.MaxPool1d(2)(out) # b*2C*64
        out = self.down2(out) # b*4C*32
        out = self.down3(out) # b*8C*16
        # out = self.down4(out) # b*16*8
        out = torch.nn.MaxPool1d(2)(out) # b*8C*8
        out = self.pread_layer(out.view(bs, -1)) # b*1024
        return out

    def init(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight)

class ResBlock(nn.Module):
    def __init__(self, n_in, n_out, stride=1):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv1d(n_in, n_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm1d(n_out)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(n_out, n_out, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(n_out)

        if stride != 1 or n_out != n_in:
            self.shortcut = nn.Sequential(
                nn.Conv1d(n_in, n_out, kernel_size=1, stride=stride),
                nn.BatchNorm1d(n_out))
        else:
            self.shortcut = None

    def forward(self, x):
        residual = x
        if self.shortcut is not None:
            residual = self.shortcut(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        
        out += residual
        out = self.relu(out)
        return out

if __name__ == "__main__":
    from torchline.config import get_cfg
    cfg = get_cfg()
    cfg.merge_from_file('../config/config.yaml')
    model = SimpleCNN(cfg)
    x = torch.rand(2,2,512)
    y = model(x)
    print(x.shape)